package com.github.manx98.okdownloader.utils;

import com.github.manx98.okdownloader.exception.ProgressDataBrokenException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.util.*;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.LongAdder;
import java.util.stream.Collectors;

/**
 * 以文件的方式保存进度
 *
 * @author : manx98
 * @date : 2021/5/26
 */
public class FileProgressHandler implements ProgressHandler {
    private static final int HEADER_SIZE = 12;
    private static final int BLOCK_SIZE = 16;
    private final RandomAccessFile progressFile;
    private final FileChannel fileChannel;
    private final Logger logger = LoggerFactory.getLogger(FileProgressHandler.class);
    private final ThreadLocal<ByteBuffer> byteBufferThreadLocal;
    private final int maxWorkers;
    private final long minSplitSize;
    private final long maxSplitSize;
    private List<DownloadBlockInfo> progressList;
    private long completedSize = 0;
    private long fileSize = -1;
    private volatile boolean closed = false;

    /**
     * 使用文件方式存储进度
     *
     * @param file         进度文件
     * @param size         下载文件大小
     * @param minSplitSize 最小拆分大小
     * @param maxSplitSize 最大拆分大小
     */
    public FileProgressHandler(File file, long size, int maxWorkers, long minSplitSize, long maxSplitSize) throws IOException {
        this.progressFile = new RandomAccessFile(file, "rw");
        this.fileChannel = this.progressFile.getChannel();
        this.fileSize = size;
        this.maxWorkers = maxWorkers;
        this.minSplitSize = minSplitSize;
        this.maxSplitSize = maxSplitSize;
        this.byteBufferThreadLocal = ThreadLocal.withInitial(() -> ByteBuffer.allocateDirect(16));
        try {
            if (this.progressFile.length() > 0) {
                loadData();
            }
            init();
        } catch (ProgressDataBrokenException e) {
            logger.warn("进度数据损坏!", e);
            this.fileSize = -1;
            this.progressFile.setLength(0);
        }
    }

    /**
     * 返回下载文件大小
     *
     * @return 如果进度文件不存在返回 -1
     */
    @Override
    public long getFileSize() {
        return fileSize;
    }

    /**
     * 获取已完成大小
     *
     * @return 如果进度文件不存在默认返回 0
     */
    @Override
    public long getCompletedSize() {
        return completedSize;
    }

    /**
     * 加载磁盘数据
     *
     * @throws IOException                 IO异常
     * @throws ProgressDataBrokenException 数据损坏
     */
    public void loadData() throws IOException, ProgressDataBrokenException {
        try {
            progressFile.seek(0);
            int blockCount = progressFile.readInt();
            this.fileSize = progressFile.readLong();
            if (blockCount < 0 || fileSize < 0) {
                throw new ProgressDataBrokenException("数据损坏");
            }
            this.progressList = new CopyOnWriteArrayList<>();
            long seek = HEADER_SIZE;
            for (int i = 0; i < blockCount; ++i, seek += BLOCK_SIZE) {
                long start = progressFile.readLong();
                long end = progressFile.readLong();
                if (start <= end) {
                    DownloadBlockInfo downloadBlockInfo = new DownloadBlockInfo(start, end, seek);
                    this.progressList.add(downloadBlockInfo);
                }
            }
            this.completedSize = this.fileSize - fixProgress(this.progressList);
            if (this.completedSize > this.fileSize || this.completedSize < 0) {
                throw new ProgressDataBrokenException("进度数据异常");
            }
        } catch (EOFException e) {
            throw new ProgressDataBrokenException("进度数据读取异常");
        }
    }

    /**
     * 加载进度
     * 进度的存储格式
     * 4    |    8     | (8,8)
     * 文件块数 |  文件大小  |文件块数据：[[起始,结束],[起始,结束],.....]
     *
     * @throws IOException IO错误
     */
    private void init() throws IOException {
        try {
            if (progressFile.length() > 0) {
                rePlanDataFragmentation(fileSize, maxWorkers, minSplitSize, maxSplitSize);
                return;
            }
        } catch (ProgressDataBrokenException e) {
            logger.warn("进度数据损坏,将被抛弃！", e);
        }
        createProgress(fileSize, maxWorkers, minSplitSize, maxSplitSize);
    }

    /**
     * 获取下载块列表
     *
     * @return 下载块列表队列
     */
    @Override
    public List<DownloadBlockInfo> getDownloadBlockProgressInfos() {
        return progressList;
    }

    @Override
    public long getMinSplitSize() {
        return minSplitSize;
    }

    @Override
    public long getMaxSplitSize() {
        return maxSplitSize;
    }

    @Override
    public int getMaxWorkers() {
        return maxWorkers;
    }

    /**
     * 重新规划数据分片
     *
     * @throws IOException                 IO异常
     * @throws ProgressDataBrokenException 进度文件损坏
     */
    private void rePlanDataFragmentation(long realFileSize, int maxWorkers, long minSplitSize, long maxSplitSize) throws IOException, ProgressDataBrokenException {
        if (this.fileSize != realFileSize) {
            throw new ProgressDataBrokenException("进度文件记录的文件长度为:" + realFileSize + "byte,与现在的文件长度不符合:" + fileSize + "byte");
        }
        if (this.completedSize > realFileSize || this.completedSize < 0) {
            throw new ProgressDataBrokenException("进度数据损坏！");
        }
        extendedThread(maxWorkers, minSplitSize, maxSplitSize, this.progressList);
        flushAll(this.progressList, fileSize);
    }

    /**
     * 将所有进度信息全部写入磁盘
     *
     * @param downloadBlockInfos 进度列表
     * @throws IOException IO异常
     */
    private void flushAll(List<DownloadBlockInfo> downloadBlockInfos, long fileSize) throws IOException {
        long seek = HEADER_SIZE;
        ByteBuffer data = ByteBuffer.allocate((downloadBlockInfos.size() << 4) + HEADER_SIZE);
        data.putInt(downloadBlockInfos.size());
        data.putLong(fileSize);
        for (DownloadBlockInfo downloadBlockInfo : downloadBlockInfos) {
            data.putLong(downloadBlockInfo.getStart());
            data.putLong(downloadBlockInfo.getEnd());
            downloadBlockInfo.setSeek(seek);
            seek += 16;
        }
        write(0, data);
    }

    /**
     * 修复下载进度信息
     *
     * @param downloadBlockInfos 下载块列表
     * @return 返回未下载完成的下载块
     */
    private long fixProgress(List<DownloadBlockInfo> downloadBlockInfos) {
        Set<DownloadBlockInfo> selected = new HashSet<>();
        for (DownloadBlockInfo downloadBlockInfo : downloadBlockInfos) {
            if (downloadBlockInfo.getStart() > downloadBlockInfo.getEnd()) {
                selected.add(downloadBlockInfo);
                continue;
            }
            if (!selected.contains(downloadBlockInfo)) {
                for (DownloadBlockInfo blockProgressInfo : downloadBlockInfos) {
                    if (downloadBlockInfo != blockProgressInfo && !selected.contains(blockProgressInfo) && downloadBlockInfo.getStart() <= blockProgressInfo.getStart() && downloadBlockInfo.getEnd() >= blockProgressInfo.getEnd()) {
                        selected.add(blockProgressInfo);
                    }
                }
            }
        }
        LongAdder uncompletedSize = new LongAdder();
        final List<DownloadBlockInfo> collect = downloadBlockInfos.stream().filter(x -> {
            boolean result = !selected.contains(x);
            if (result) {
                uncompletedSize.add(x.getEnd() - x.getStart() + 1);
            }
            return result;
        }).collect(Collectors.toList());
        downloadBlockInfos.clear();
        downloadBlockInfos.addAll(collect);
        return uncompletedSize.longValue();
    }

    /**
     * 尝试再次划分下载块，使总分片数量达到最大线程数
     *
     * @param maxWorker    下载线程数
     * @param minSplitSize 最小分割长度
     */
    private void extendedThread(int maxWorker, long minSplitSize, long maxSplitSize, List<DownloadBlockInfo> downloadBlockInfos) {
        for (int i = downloadBlockInfos.size(); i < maxWorker; ++i) {
            if (!extendTaskCount(minSplitSize, maxSplitSize, downloadBlockInfos)) {
                break;
            }
        }
    }

    /**
     * 再次划分下载任务，增加一个下载块
     *
     * @param minSplitSize 最小分割
     * @return 拓展是否成功
     */
    private boolean extendTaskCount(long minSplitSize, long maxSplitSize, List<DownloadBlockInfo> downloadBlockInfos) {
        DownloadBlockInfo longerest = null;
        long minSplitSizeX2 = minSplitSize << 1;
        for (DownloadBlockInfo progressInfo : downloadBlockInfos) {
            long len = progressInfo.getStart() - progressInfo.getEnd() + 1;
            if (len >= minSplitSizeX2 || len > maxSplitSize) {
                if (longerest == null) {
                    longerest = progressInfo;
                } else {
                    if (longerest.getEnd() - longerest.getStart() + 1 < len) {
                        longerest = progressInfo;
                    }
                }
            }
        }
        if (longerest != null) {
            long start = longerest.getStart();
            long end = longerest.getEnd();
            long mid = start + ((start - end + 1) >> 1);
            if (mid + 1 <= end) {
                longerest.setEnd(mid);
                downloadBlockInfos.add(new DownloadBlockInfo(mid + 1, end, 0));
                return true;
            }
        }
        return false;
    }

    /**
     * 写入数据
     *
     * @param seek 文件下标
     * @param data 数据
     * @throws IOException IO异常
     */
    private void write(long seek, ByteBuffer data) throws IOException {
        data.flip();
        try {
            fileChannel.write(data, seek);
        } finally {
            data.clear();
        }
    }

    /**
     * 创建文件切片列表
     *
     * @param size         文件大小
     * @param maxWorkers   最大下载线程
     * @param minSplitSize 文件切片最小值
     * @throws IOException IO异常
     */
    private void createProgress(long size, int maxWorkers, long minSplitSize, long maxSplitSize) throws IOException {
        long splitSize = Math.min(Math.max(minSplitSize, size / maxWorkers), maxSplitSize);
        long start = 0;
        progressFile.setLength(0);
        List<DownloadBlockInfo> downloadBlockInfos = new CopyOnWriteArrayList<>();
        long seek = 0;
        splitSize -= 1;
        long lastSeek = size - 1;
        while (start <= lastSeek) {
            long end = Math.min(lastSeek, start + splitSize);
            downloadBlockInfos.add(new DownloadBlockInfo(start, end, seek));
            start = end + 1;
            seek += 16;
        }
        flushAll(downloadBlockInfos, size);
        this.progressList = downloadBlockInfos;
        this.fileSize = size;
        this.completedSize = 0;
    }

    @Override
    public void updateStartProgress(DownloadBlockInfo progressInfo) throws IOException {
        final ByteBuffer buffer = byteBufferThreadLocal.get();
        buffer.putLong(progressInfo.getStart());
        write(progressInfo.getSeek(), buffer);
    }

    @Override
    public void updateEndProgress(DownloadBlockInfo progressInfo) throws IOException {
        final ByteBuffer buffer = byteBufferThreadLocal.get();
        buffer.putLong(progressInfo.getEnd());
        write(progressInfo.getSeek() + 8, buffer);
    }

    @Override
    public void updateProgress(DownloadBlockInfo progressInfo) throws IOException {
        final ByteBuffer buffer = byteBufferThreadLocal.get();
        buffer.putLong(progressInfo.getStart());
        buffer.putLong(progressInfo.getEnd());
        write(progressInfo.getSeek(), buffer);
    }

    @Override
    public void onThreadTermination() {
        byteBufferThreadLocal.remove();
    }

    /**
     * 资源是否已经关闭
     */
    public boolean isClosed() {
        return closed;
    }

    @Override
    public void close() throws Exception {
        if (!closed) {
            synchronized (this) {
                if (!closed) {
                    fileChannel.close();
                    progressFile.close();
                    closed = true;
                }
            }
        }
    }
}
